-- | Module:        Hexena.Chunks
--   Stability:     Experimental
--   Portability:   Portable
--  
-- <TODO>
module Hexena.Chunks
  ( Chunkable
  , IFFMap
  , Testtype
  , bodyID
  , isForm
  , isFormId
  , loadBODY
  , loadMalwIFF
  , loadMalwIFF'
  , loadMalwareFile
  , loadANNO
  , saveANNO
  -- , saveBLEN
  , loadFORMHIST
  , makeFORMHIST
  , mapIFF
  , loadPEHD
  , savePEHD
  , saveHASH
  , malwId
  , writeIFF
  ) where

import qualified Codec.Compression.BZip as CBZ
import qualified Codec.Compression.GZip as CGZ
import Codec.Utils (Octet)
import Control.Monad (ap)
import Data.Array.Unboxed (elems, listArray)
import Data.Binary
import Data.Binary.Get
import Data.Binary.Put
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import Data.Char (chr, ord)
import qualified Data.Digest.SHA1 as SHA1
import Data.IFF
import Data.List (find)
import qualified Data.Map as M
import Data.Maybe

import Hexena.PEHeader
import Hexena.Math.Statistics

data Testtype = Testtype Word32 Word16
type IFFMap = M.Map String Chunk

bodyID :: ChunkId
bodyID = chunkIdFromString "BODY"

isFormId :: ChunkId -> T -> Bool
isFormId cid (Cons c (Form _)) = c == cid
isFormId _ _ = False

isForm :: T -> Bool
isForm (Cons _ (Form _)) = True
isForm _ = False

-- | Load a malware file and return either a hexena file of just a
-- bytestring.
loadMalwareFile :: String -> IO (Either BL.ByteString T)
loadMalwareFile fn = let pIFF Nothing = error "File is not an IFF file!"
                         pIFF (Just x) = x
                     in do x <- B.readFile fn
                           if [0x4d, 0x5a] == take 2 (B.unpack x) then return (Left $ BL.fromChunks [x]) else return (Right $ pIFF $ fromByteString x)

-- | Load a MALW IFF file. It will throw an exception if the file can
-- | not be loaded, is not an IFF file or if it is not a FORM.MALW IFF
-- | file.
loadMalwIFF :: String -> IO IFFMap
loadMalwIFF fname = do dat <- fromByteString `fmap` B.readFile fname
                       case dat of
                         Nothing -> error "File is not an IFF file!"
                         Just x -> return $ mapIFF x

-- | Load a MALW IFF file. It will throw an exception if the file can
-- | not be loaded, is not an IFF file or if it is not a FORM.MALW IFF
-- | file.
loadMalwIFF' :: String -> IO (IFFMap, [T])
loadMalwIFF' fname = let deiff (Cons cid (Form xs)) = xs
                     in do dat <- fromByteString `fmap` B.readFile fname
                           case dat of
                             Nothing -> error "File is not an IFF file!"
                             Just x -> return (mapIFF x, deiff x)


mapIFF :: T -> IFFMap
mapIFF (Cons cid (Form xs)) | cid == malwId = M.fromList $ map f xs 
                            | otherwise = error "Guru Meditation! File is not an MALW IFF file!"
    where f (Cons i c) = (chunkIdToString i, c)
mapIFF _ = error "Guru Meditation! File is not a FORM IFF file!"

decompressBodyF :: Word8 -> B.ByteString -> B.ByteString
decompressBodyF 0 = id
decompressBodyF 2 = \x -> B.concat $ BL.toChunks $ CGZ.decompress $ BL.fromChunks [x]
decompressBodyF 3 = \x -> B.concat $ BL.toChunks $ CBZ.decompress $ BL.fromChunks [x]
decompressBodyF _ = error "Illegal compression level"

loadBODY :: [T] -> B.ByteString
loadBODY iff = let body = content $ chunk $ fromJust $ find ((bodyID ==) . chunkId) iff
                   mwhd = content $ chunk $ fromJust $ find ((chunkIdFromString "MWHD" ==) . chunkId) iff
                   comp = runGet (getWord32be >> getWord8) $ BL.fromChunks [mwhd]
               in decompressBodyF comp body

malwId :: ChunkId
malwId = chunkIdFromString "MALW"

-- | Class for managing chunkable data. This will make it easiser to
-- | handle marshalling of data, etc.
class Chunkable a where
    chunkify :: a -> T
    chunkify x = Cons (chunkifyId x) $ chunkifyData x
    chunkifyId :: a -> ChunkId
    chunkifyData :: a -> Chunk
    -- How to do this?
    -- unchunkify :: T -> Maybe a
    -- Fails with "Chunks.hs:24:51: Not in scope: `a'":
    -- unchunkify (Cons id ch) = if id == chunkableId a then unchunkifyChunk ch else Nothing
    -- unchunkifyChunk :: Chunk -> a
    -- ifOfType :: T -> Bool



instance Chunkable Testtype where
    chunkifyId (Testtype _ _) = chunkIdFromString "TSTY"
    chunkifyData (Testtype w32 w16) = Chunk $ B.concat $ BL.toChunks $ runPut $ do
                                        putWord32be w32
                                        putWord16be w16

makeFORMHISTHIST :: Histogram -> String -> T
makeFORMHISTHIST h s = Cons (chunkIdFromString s) $ Chunk $ B.concat $ BL.toChunks $ runPut $ putF $ elems h
    where putF = foldr ((>>) . putWord32be . fromIntegral) (return ())

makeFORMHIST :: Histogram -> Histogram -> T
makeFORMHIST h r = Cons (chunkIdFromString "HIST") $ Form [makeFORMHISTHIST h "1GRM", makeFORMHISTHIST r "1RGR"]

-- Format for histogram saving:
--
-- FORM.MALW/FORM.HIST: Container for all histogram related stuff
-- FORM.MALW/FORM.HIST/1GRM: normal 1-gram histogram (1D)
-- FORM.MALW/FORM.HIST/1RGR: run-length reduced 1-gram histogram (1D)

-- | Load the histograms from the IFF file. First the normal, second
-- | the rle reduced histogram is returned.
loadFORMHIST :: IFFMap -> (Histogram, Histogram)
loadFORMHIST iff = (lA chunk1grm, lA chunk1rgr)
    where chks = parts $ fromJust $ M.lookup "HIST" iff
          find' n = BL.fromChunks [content $ chunk $ fromJust $ find ((==) (chunkIdFromString n) . chunkId) chks]
          chunk1grm = find' "1GRM"
          chunk1rgr = find' "1RGR"
          lA = listArray (0, 255) . map fromIntegral . runGet (getF 256)
          getF :: Int -> Get [Word32]
          getF 0 = return []
          getF x = getWord32be >>= (\y -> (y :) `fmap` getF (x - 1))


-- PE Headers are just the PE header in binary representation
-- | save a PEHD chunk from a PEHeader
savePEHD :: PEHeader -> T
savePEHD pe = Cons (chunkIdFromString "PEHD") $ Chunk $ B.concat $ BL.toChunks $ runPut putF
  where peoh = peOptionalHeader pe
        putF = do  
          putWord32be $ peSignature pe
          putWord16be $ peMachine pe
          putWord16be $ peNumberOfSections pe
          putWord32be $ peTimeDateStamp pe
          putWord32be $ pePointerToSymbolTable pe
          putWord32be $ peNumberOfSymbols pe
          putWord16be $ peSizeOfOptionalHeader pe
          putWord16be $ peCharacteristics pe
          -- PE optional header from here on...
          putWord16be $ peohMagic peoh
          putWord8 $ peohMajorLinkerVersion peoh
          putWord8 $ peohMinorLinkerVersion peoh
          putWord32be $ peohSizeOfCode peoh
          putWord32be $ peohSizeofInitialisedData peoh
          putWord32be $ peohSizeofUninitialisedData peoh
          putWord32be $ peohAddressOfEntryPoint peoh
          putWord32be $ peohBaseOfCode peoh
          putWord32be $ peohBaseOfData peoh
          putWord32be $ peohImageBase peoh
          putWord32be $ peohSectionAlignment peoh
          putWord32be $ peohFileAlignment peoh
          putWord16be $ peohMajorOSVersion peoh
          putWord16be $ peohMinorOSVersion peoh
          putWord16be $ peohMajorImageVersion peoh
          putWord16be $ peohMinorImageVersion peoh
          putWord16be $ peohMajorSubsystemVersion peoh
          putWord16be $ peohMinorSubsystemVersion peoh
          putWord32be $ peohReserved peoh
          putWord32be $ peohSizeOfImage peoh
          putWord32be $ peohSizeOfHeaders peoh
          putWord32be $ peohCheckSum peoh

-- | load a PEHeader from a Chunk
loadPEHD :: IFFMap -> PEHeader
loadPEHD iff = runGet load chk
    where chk = BL.fromChunks [content $ fromJust $ M.lookup "PEHD" iff]
          load = return PEHeader `ap` getWord32be
                                 `ap` getWord16be
                                 `ap` getWord16be
                                 `ap` getWord32be
                                 `ap` getWord32be
                                 `ap` getWord32be
                                 `ap` getWord16be
                                 `ap` getWord16be
                                 `ap` (PEOptionalHeader `fmap` getWord16be
                                       `ap` getWord8
                                       `ap` getWord8
                                       `ap` getWord32be
                                       `ap` getWord32be
                                       `ap` getWord32be
                                       `ap` getWord32be
                                       `ap` getWord32be
                                       `ap` getWord32be
                                       `ap` getWord32be
                                       `ap` getWord32be
                                       `ap` getWord32be
                                       `ap` getWord16be
                                       `ap` getWord16be
                                       `ap` getWord16be
                                       `ap` getWord16be
                                       `ap` getWord16be
                                       `ap` getWord16be
                                       `ap` getWord32be
                                       `ap` getWord32be
                                       `ap` getWord32be
                                       `ap` getWord32be)

-- | Annotation for a file.

-- According to the IFF FORM and Chunk registry an annotation chunk is
-- a generic chunk used for file annotations. We use this one here too
-- (and it will be put into the report).
loadANNO :: IFFMap -> String
loadANNO = f . M.lookup "ANNO"
    where f Nothing = []
          f (Just x) = map (chr . fromIntegral) $ B.unpack $ content x

-- | Annotation chunk (IFF standard)
saveANNO :: String -> T
saveANNO = Cons (chunkIdFromString "ANNO") . Chunk . B.pack . map (fromIntegral . ord)


-- | Save the standard hashes (currently MD5 and SHA1, SHA256) to an IFF chunk.
saveHASH :: [Word8] -> SHA1.Word160 -> [Octet] -> T
saveHASH m (SHA1.Word160 x1 x2 x3 x4 x5) s256 =  Cons (chunkIdFromString "HASH") $ Form [mkm, mks, mks256]
  where mkm = Cons (chunkIdFromString "MD5 ") $ Chunk $ B.pack m
        mks = Cons (chunkIdFromString "SHA1") $ Chunk $ B.concat $ BL.toChunks $ runPut pf
        mks256 = Cons (chunkIdFromString "S256") $ Chunk $ B.pack s256
        pf = do putWord32be x1
                putWord32be x2
                putWord32be x3
                putWord32be x4
                putWord32be x5

writeIFF :: String -> [T] -> IO ()
writeIFF fn = B.writeFile fn . toByteString . Cons (chunkIdFromString "MALW") . Form


-- 
